/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.extensions.sql.meta.provider.kafka;



import com.bff.gaia.unified.sdk.schemas.Schema;

import com.bff.gaia.unified.sdk.transforms.DoFn;

import com.bff.gaia.unified.sdk.transforms.PTransform;

import com.bff.gaia.unified.sdk.transforms.ParDo;

import com.bff.gaia.unified.sdk.values.KV;

import com.bff.gaia.unified.sdk.values.PCollection;

import com.bff.gaia.unified.sdk.values.Row;

import org.apache.commons.csv.CSVFormat;

import com.bff.gaia.unified.sdk.extensions.sql.impl.schema.UnifiedTableUtils;



import java.util.List;



import static java.nio.charset.StandardCharsets.UTF_8;



/** A Kafka topic that saves records as CSV format. */

public class UnifiedKafkaCSVTable extends UnifiedKafkaTable {

  private CSVFormat csvFormat;



  public UnifiedKafkaCSVTable(Schema unifiedSchema, String bootstrapServers, List<String> topics) {

    this(unifiedSchema, bootstrapServers, topics, CSVFormat.DEFAULT);

  }



  public UnifiedKafkaCSVTable(

	  Schema unifiedSchema, String bootstrapServers, List<String> topics, CSVFormat format) {

    super(unifiedSchema, bootstrapServers, topics);

    this.csvFormat = format;

  }



  @Override

  public PTransform<PCollection<KV<byte[], byte[]>>, PCollection<Row>> getPTransformForInput() {

    return new CsvRecorderDecoder(schema, csvFormat);

  }



  @Override

  public PTransform<PCollection<Row>, PCollection<KV<byte[], byte[]>>> getPTransformForOutput() {

    return new CsvRecorderEncoder(schema, csvFormat);

  }



  /** A PTransform to convert {@code KV<byte[], byte[]>} to {@link Row}. */

  public static class CsvRecorderDecoder

      extends PTransform<PCollection<KV<byte[], byte[]>>, PCollection<Row>> {

    private Schema schema;

    private CSVFormat format;



    public CsvRecorderDecoder(Schema schema, CSVFormat format) {

      this.schema = schema;

      this.format = format;

    }



    @Override

    public PCollection<Row> expand(PCollection<KV<byte[], byte[]>> input) {

      return input

          .apply(

              "decodeRecord",

              ParDo.of(

                  new DoFn<KV<byte[], byte[]>, Row>() {

                    @ProcessElement

                    public void processElement(ProcessContext c) {

                      String rowInString = new String(c.element().getValue(), UTF_8);

                      for (Row row : UnifiedTableUtils.csvLines2UnifiedRows(format, rowInString, schema)) {

                        c.output(row);

                      }

                    }

                  }))

          .setRowSchema(schema);

    }

  }



  /** A PTransform to convert {@link Row} to {@code KV<byte[], byte[]>}. */

  public static class CsvRecorderEncoder

      extends PTransform<PCollection<Row>, PCollection<KV<byte[], byte[]>>> {

    private Schema schema;

    private CSVFormat format;



    public CsvRecorderEncoder(Schema schema, CSVFormat format) {

      this.schema = schema;

      this.format = format;

    }



    @Override

    public PCollection<KV<byte[], byte[]>> expand(PCollection<Row> input) {

      return input.apply(

          "encodeRecord",

          ParDo.of(

              new DoFn<Row, KV<byte[], byte[]>>() {

                @ProcessElement

                public void processElement(ProcessContext c) {

                  Row in = c.element();

                  c.output(KV.of(new byte[] {}, UnifiedTableUtils.unifiedRow2CsvLine(in, format).getBytes(UTF_8)));

                }

              }));

    }

  }

}